﻿using System;
using System.Diagnostics;
using RayDen.Library.Core;
using RayDen.Library.Core.Primitives;
using RayDen.RayEngine.Data;

namespace RayDen.RayEngine.VertexFramework
{

    public enum GenericPathSamplerState
    {
        Initialized,
        Sampling,
        Connection,
        Evaluation,
    }

    public struct PathConnection
    {
        public RayData Ray;
        public RgbSpectrum Throughput;
        public float Pdf;
        public bool Connected;
        public int PrimitiveIndex, CurrentShadowRayIndex;
        public int VertexIndex;
    }

    public class GenericPathSampler : VertexPathSamplerBase
    {
        public PathElement[] vertices;
        public PathConnection[] connections;

        public int currentVertice;
        internal RgbSpectrum Throughput;
        public GenericPathSamplerState PathState;

        protected int depth, tracedShadowRayCount, indexRef, verticesCount;
        protected float pathWeight;
        protected bool specularBounce;

        protected internal RayEngineScene scene;

        public GenericPathSampler()
        {
            vertices = new PathElement[VerticesMax];
        }

        public int MaxRaysPerPath
        {
            get;
            set;
        }

        public override void InitPath(PathBufferBase buffer)
        {
            base.InitPath(buffer);
            this.scene = buffer.Scene;
            this.Radiance = new RgbSpectrum(0f);
            this.Throughput = new RgbSpectrum(1f);
            this.PathState = GenericPathSamplerState.Initialized;
            //vertices = new PathVertex[scene.MaxPathDepth * 2];
            this.Sample = buffer.Sampler.GetSample();
            IRay ray;
            buffer.Scene.Camera.GetRay(Sample.imageX, Sample.imageY, out ray);
            currentVertice = 0;
            vertices[currentVertice] = new ObserverPathElement() { CameraSample = new CameraSample() { EyeRay = (RayData)ray, imageX = Sample.imageX, imageY = Sample.imageY } };
            currentVertice++;
            this.PathRay = (RayData)ray;
            this.RayIndex = -1;
            this.pathWeight = 1.0f;
            this.depth = 0;
            this.verticesCount = 0;
            this.specularBounce = true;
        }

        public override bool FillRayBuffer(RayBuffer rayBuffer)
        {
            var leftSpace = rayBuffer.LeftSpace();
            if (((this.PathState == GenericPathSamplerState.Initialized) && (1 > leftSpace)) ||
                ((PathState == GenericPathSamplerState.Connection) && (tracedShadowRayCount > leftSpace)) ||
                ((this.PathState == GenericPathSamplerState.Sampling) && (1 > leftSpace)))
                return false;
            if (PathState == GenericPathSamplerState.Connection)
            {
                for (int i = 0; i < tracedShadowRayCount; ++i)
                    connections[i].CurrentShadowRayIndex = rayBuffer.AddRay(connections[i].Ray);
            }
            else
            {
                RayIndex = rayBuffer.AddRay(PathRay);
            }
            return true;
        }

        public override void Advance(RayBuffer rayBuffer, SampleBuffer consumer, Action<PathSamplerBase> onRestart)
        {
            int currentTriangleIndex = 0;
            SurfaceIntersectionData hitInfo = null;

#if VERBOSE

            try
            {
#endif

                base.Advance(rayBuffer, consumer, onRestart);



                if (((this.PathState == GenericPathSamplerState.Connection)))
                {
                    for (int i = 0; i < tracedShadowRayCount; i++)
                    {
                        connections[i].Connected = rayBuffer.rayHits[connections[i].CurrentShadowRayIndex].Miss();
                    }

                    PathState = GenericPathSamplerState.Evaluation;
                    return;
                }
                var rayHit = rayBuffer.rayHits[RayIndex];

                depth++;
                bool missed = rayHit.Index == 0xffffffffu;


                if (missed)
                {
                    if (depth <= scene.MaxPathDepth)
                    {
                        vertices[currentVertice] = new EnvironmentPathElement()
                            {
                                Throughput = SampleEnvironment(-PathRay.Dir)
                            };
                    }
                    if (PathState == GenericPathSamplerState.Initialized || PathState == GenericPathSamplerState.Sampling)
                    {
                        PathState = GenericPathSamplerState.Connection;
                        return;
                    }
                }

                // Something was hit
                hitInfo = SurfaceSampler.GetIntersection(ref PathRay, ref rayHit);

                currentTriangleIndex = (int)rayHit.Index;


                if (hitInfo == null)
                {
                    Debugger.Break();
                    return;
                }
                var hitPoint = PathRay.Point(rayHit.Distance);

                //If Hit light)
                if (hitInfo.IsLight)
                {
                    var lt = scene.GetLightByIndex(currentTriangleIndex);
                    vertices[currentVertice] = new LightsourcePathElement()
                        {
                            HitInfo = hitInfo,
                            Throughput = this.Throughput,
                            LightSample = new LightSample()
                                {
                                    Spectrum = lt.Le(ref PathRay.Dir),
                                    Pdf = 1f

                                }
                        };
                    PathState = GenericPathSamplerState.Evaluation;
                    return;
                }


                Vector wo = -PathRay.Dir;



                float fPdf;
                Vector wi;
                float u0 = Sample.GetLazyValue(), u1 = Sample.GetLazyValue(), u2 = Sample.GetLazyValue();
                RgbSpectrum f = hitInfo.Material.Sample_f(ref wo, out wi, ref hitInfo.Normal, ref hitInfo.ShadingNormal,
                                                          u0, u1, u2,
                                                          out fPdf, out specularBounce) * hitInfo.Color;

                vertices[currentVertice] = new GeometryPathElement()
                    {
                        HitInfo = hitInfo,
                        HitPoint = hitPoint,
                        BsdfSample = new BsdfSample()
                            {
                                SampleData = new[] { u0, u1, u2 },
                                Wi = wi,
                                Pdf = fPdf,
                                Spectrum = f.ToArray(),
                                SpecularBounce = specularBounce
                            }
                    };
                if (hitInfo.Material.IsDiffuse())
                {
                    float lightStrategyPdf = scene.ShadowRayCount / (float)scene.Lights.Length;
                    RgbSpectrum lightTroughtput = Throughput * hitInfo.Color;
                    for (int i = 0; i < scene.ShadowRayCount; ++i)
                    {
                        int currentLightIndex = scene.SampleLights(Sample.GetLazyValue());
                        var light = scene.Lights[currentLightIndex];

                        var ls = new LightSample();
                        light.EvaluateShadow(ref hitPoint, ref hitInfo.ShadingNormal, Sample.GetLazyValue(), Sample.GetLazyValue(), Sample.GetLazyValue(), out ls);
                        if (ls.Pdf <= 0f)
                        {
                            continue;
                        }

                        connections[tracedShadowRayCount].Throughput = new RgbSpectrum(ls.Spectrum);
                        connections[tracedShadowRayCount].Pdf = ls.Pdf;
                        connections[tracedShadowRayCount].Ray = ls.LightRay;

                        Vector lwi = connections[tracedShadowRayCount].Ray.Dir;
                        connections[tracedShadowRayCount].Throughput *= lightTroughtput *
                                                               Vector.AbsDot(ref hitInfo.ShadingNormal, ref lwi) *
                                                               hitInfo.Material.f(
                                                                   ref connections[tracedShadowRayCount].Ray.Dir,
                                                                   ref wo, ref hitInfo.ShadingNormal);
                        if (!connections[tracedShadowRayCount].Throughput.IsBlack())
                        {
                            connections[tracedShadowRayCount].Pdf *= lightStrategyPdf;
                            tracedShadowRayCount++;
                        }
                    }
                }

                if ((fPdf <= 0.0f) || f.IsBlack())
                {
                    PathState = GenericPathSamplerState.Connection;
                    return;
                }
                pathWeight *= fPdf;
                Throughput *= f / fPdf;


                if (depth > scene.MaxPathDepth)
                {
                    float prob = Math.Max(Throughput.Filter(), scene.RussianRuletteImportanceCap);
                    if (prob >= Sample.GetLazyValue())
                    {

                        Throughput /= prob;
                        pathWeight *= prob;
                    }
                    else
                    {
                        PathState = GenericPathSamplerState.Connection;
                        return;
                    }
                }

                PathRay.Org = hitPoint;
                PathRay.Dir = wi.Normalize();
                this.PathState = GenericPathSamplerState.Sampling;
#if VERBOSE
            }
            catch (Exception ex)
            {
                Console.WriteLine("Advance path exception");
                Console.WriteLine("Error triangle {0}", currentTriangleIndex);
                Console.WriteLine(ex.Message);
                Console.WriteLine(ex.StackTrace);

            }
#endif
        }

        private RgbSpectrum SampleEnvironment(Vector vector)
        {
            return this.scene.SampleEnvironment(vector);
        }
    }
}